#!/usr/local/bin/python
#-*- encoding:utf-8 -*-

import numpy as np;
import scipy.sparse as sp
from scipy import linalg
import time
import threading
import os
from scipy.sparse.linalg import eigsh
from multiprocessing import Process, Manager
import sys
import utils


config = utils.get_config();
beta = config.getfloat('rmls', 'beta');
gamma = config.getfloat('rmls', 'gamma')
thetax = config.getfloat('rmls', 'thetax');
thetay = config.getfloat('rmls', 'thetay');
converge_threshhold= config.getfloat('rmls', 'converge_threshhold');
dx = config.getint('rmls', 'dx');
dy = config.getint('rmls', 'dy');
d = config.getint('rmls', 'd');

def saveModel(Lx, Ly, batch):
    print 'Saving Model For Batch %d' % batch;
    LxFile = 'Lx_%d.txt' % batch;
    LyFile = 'Ly_%d.txt' % batch;

    utils.save_sparse_csr(LxFile, Lx.tocsr());
    utils.save_sparse_csr(LyFile, Ly.tocsr());
   
def calcC(row, theta):
    t = np.linalg.norm(row);
    if t == 0:
        return 0;
    return theta/t;

def checkConverge(LPx,LPy, Lx, Ly):
    t = np.linalg.norm((Lx-LPx).toarray()) + np.linalg.norm((Ly-LPy).toarray());
    if t < converge_threshhold:
        return (True, t);
    else:
        return (False, t);

def checkConvergeM(LPx, LPy, Lx, Ly):
    t = utils.sparseFroNorm(LPx - Lx) + utils.sparseFroNorm(LPy - Ly);
    if t < converge_threshhold:
        return (True, t);
    else:
        return (False, t);

def compute(x, t):
    return np.sign(x) * max(np.abs(x)-t, 0);

def updateLxu(slx, Lx, pid, start, end ,debug=False):
    #print 'pid: %d, start: %d, end: %d' % (pid, start, end);
    compute_lxu = np.frompyfunc(lambda x: compute(x, beta), 1, 1);
    for u in range(start, end):
        if debug:
            start = time.time();
        wu = WX.getcol(u);

        lxu= compute_lxu(wu.transpose().todense());

        #cu = calcC(lxu, thetax);
        #Lx[u] =  lxu * cu;

        if debug:
            print  u, time.time() -start;
    slx.append((Lx, start, end));

def updateLyv(sly, Ly, pid, start, end, debug=False):
    #print 'pid: %d, start: %d, end: %d' % (pid, start, end);
    compute_lyv = np.frompyfunc(lambda x:compute(x, gamma), 1, 1);
    for v in range(start, end):
        if debug:
            start = time.time();
        nv = NY.getcol(v);
        lyv = compute_lyv(nv.transpose().todense());
        
        #cv = calcC(lyv, thetay);
        #Ly[v] = lyv * cv;

        if debug:
            print v,time.time() -start;
    sly.append((Ly, start, end));

def update(M, sl):
    t = 0;
    for triple in sl:
        R, s, e = triple;
        M[s:e] = R[s:e];
        
    return t;

def RMLS(W, N, T=25, batchno=-1):
    con = -1;
    global WX, NY, Lx, Ly;
    K = 20;
    #Generate Lx and Ly;
    I = sp.identity(dx);
    evals, evecs = eigsh(I, d, which='LM', ncv=3*d);
    Lx = sp.lil_matrix(evecs);
    Ly = sp.lil_matrix(evecs);
    #Lx = sp.rand(dx, d, density=0.1, format='lil', dtype=np.float64);
    #Ly = sp.rand(dy, d, density=0.1, format='lil', dtype=np.float64);
    manager = Manager();
    t = 0
    con = utils.sparseFroNorm(Lx) + utils.sparseFroNorm(Ly);
    while True:
        t += 1;
        LPx = Lx.copy();
        LPy = Ly.copy();
        print 'This is %d iteration, last converge is %f' %(t, con);
        start = time.time();
        WX = Ly.tocsr().transpose().dot(W);
        #print 'WX', type(WX);
        ranx = dx/K;

        processListx = [];
        u = 0;
        slx = manager.list();
        for pid in range(K-1):
            processListx.append(Process(target=updateLxu, args=(slx, Lx, pid, u, u+ranx)));
            u = u + ranx;
        processListx.append(Process(target=updateLxu, args=(slx, Lx, K-1, u, dx, False)));

        for pid in range(K):
            processListx[pid].start();
        for pid in range(K):
            processListx[pid].join();
    
        norm =  update(Lx, slx);
        utils.l2norm(Lx, d);
 
        print 'Update Lx cost: ' + str(time.time() -start); 
    
        start = time.time();
        NY = Lx.tocsr().transpose().dot(N);
        #print 'NY', type(NY);
        #Update Ly
        rany = dy/K;
        processListy = [];
        v  = 0;
        sly  = manager.list();
        for pid in range(K-1):
            processListy.append(Process(target=updateLyv, args=(sly,Ly,pid, v, v+rany)));
            v = v + rany;
        processListy.append(Process(target=updateLyv, args=(sly, Ly, K-1, v, dy, False)));

        for pid in range(K):
            processListy[pid].start();
        for pid in range(K):
            processListy[pid].join();

        norm +=  update(Ly, sly);  
        utils.l2norm(Ly, d);  
        print 'Update Ly cost: ' + str(time.time() - start);

        #print LPx, Lx;
        flag, con = checkConvergeM(LPx, LPy, Lx, Ly);
       
        print 'Converge: ', flag, con;
        if flag:
            break;
        if t >= T:
            break;

    return (Lx, Ly);
   

def _max_min(max, min, x):
    return (x-min)/(max-min);
#ONLY FOR CSC
def max_min(W):
    WF = W.todense();
    max = WF.max();
    min = WF.min();
    mean = WF.mean();
    W = (WF -mean)/(max-min);
    return sp.csc_matrix(W);

if __name__ == '__main__':
    batchno = 20;
    try:
        W = utils.load_sparse_csc('W.npz');
        #N = utils.load_sparse_csc('NF.npz');
        #W = max_min(W);
        N = W.transpose().tocsc();
        print W.max(), W.min(), N.max(), N.min();
        print W;
        print N.shape;
      
        RMLS(W,N, 20, batchno); 
        saveModel(Lx, Ly, batchno);
    except KeyboardInterrupt:
        saveModel(Lx, Ly, batchno);
        sys.exit(0);
